


fen = open("train.tags.zh-en.en", encoding="utf-8")
fzh = open("train.tags.zh-en.zh", encoding="utf-8")
en_zh = []


while True:
    lz = fzh.readline()
    le = fen.readline()
    if not lz:
        assert not le
        break
    lz, le = lz.strip(), le.strip()
    
    #解析文件各部分
    if lz.startswith("<url>"):
        assert le.startswith("<url>")
        lz = fzh.readline()
        le = fen.readline()
        
        #关键词
        assert lz.startswith("<keywords>")
        assert le.startswith("<keywords>")
        lz = fzh.readline()
        le = fen.readline()
        
        #演讲人
        assert lz.startswith("<speaker>")
        assert le.startswith("<speaker>")
        lz = fzh.readline()
        le = fen.readline()
        
        #演讲ID
        assert lz.startswith("<talkid>")
        assert le.startswith("<talkid>")
        lz = fzh.readline()
        le = fen.readline()
        
        #标题
        assert lz.startswith("<title>")
        assert le.startswith("<title>")
        lz = fzh.readline()
        le = fen.readline()
        
        #描述
        assert lz.startswith("<description>")
        assert le.startswith("<description>")
        
    else:
        if not lz:
            assert not le
            break
        lee = []
        for w in le.split(" "):  #对一句话中的单词进行分隔
            w = w.replace(".","").replace(",","").lower()
            if w:
                lee.append(w)
        en_zh.append([lee, list(lz)])  #en_zh = [[[英文], [中文]], [[英文句], [中文句]], ..., [[], []]]


print(en_zh[:3])


#统计中英文词的各自的数量
from tqdm import tqdm
en_words = set()
zh_words = set()

for s in tqdm(en_zh):
    for word in s[0]:
        en_words.add(word)
    for word in s[1]:
        if word:
            zh_words.add(word)


#在中英文各自词表中添加特殊标记词
en_wl = ["<sos>", "<eos>", "<pad>"] + list(en_words)
zh_wl = ["<sos>", "<eos>", "<pad>"] + list(zh_words)

pad_id = 2
en2id = {}
zh2id = {}

for i, w in enumerate(en_wl):
    en2id[w] = i

for i, w in enumerate(zh_wl):
    zh2id[w] = i


en2id["this"]





import random
random.shuffle(en_zh)
data_len = len(en_zh)
train_data = en_zh[: int(data_len * 0.8)]
test_data = en_zh[int(data_len * 0.8):]


import torch
batch_size = 16
#data_work = 8


#构建Dataset子类
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, examples):
        self.examples = examples
        
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, index):
        example = self.examples[index]
        s1 = example[0]  #英文
        s2 = example[1]  #中文
        
        len1 = len(s1)
        len2 = len(s2)
        return s1, len1, s2, len2, index


def the_collate_fn(batch):
    batch_size = len(batch)
    src = [[0] * batch_size]
    tar = [[0] * batch_size]
    #src = []
    #tar = []
    
    src_max_1 = 0
    for b in batch:
        src_max_1 = max(src_max_1, b[1])  #一个batch中最长的英文句子的长度
    tar_max_1 = 0
    for b in batch:
        tar_max_1 = max(tar_max_1, b[3])  #一个batch中最长的中文句子的长度
    
    for x in batch:  #x: [s1, l1, s2, l2,index]
        l = []
        for w in x[0]:
            l.append(en2id[w])
        if(len(l) < src_max_1):
            for epoch in range(src_max_1 - len(l)):
                l.append(pad_id)
        src.append(l)
        
    for x in batch:
        l = []
        for w in x[2]:
            l.append(zh2id[w])
        if(len(l) < tar_max_1):
            for epoch in range(tar_max_1 - len(l)):
                l.append(pad_id)
        tar.append(l)

    index = [b[4] for b in batch]
    src.append([1] * batch_size)
    tar.append([1] * batch_size)
    s1 = torch.LongTensor(src)
    s2 = torch.LongTensor(tar)
    return s1, s2, index


#构建Dataset和DataLoader
train_dataset = MyDataset(train_data)
test_dataset = MyDataset(test_data)

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size,
    shuffle=True,
    collate_fn=the_collate_fn,
)

test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size,
    shuffle=False,
    collate_fn=the_collate_fn
)





#编码器与解码器均选择LSTM


#编码器
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, word_count, embed_dim, hidden_dim, n_layers, dropout):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers
        self.embedding = nn.Embedding(word_count, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim,n_layers, dropout=dropout)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self,src):
        embedded = self.dropout(self.embedding(src))
        outputs, (hidden, cell) = self.lstm(embedded)
        return hidden, cell


#解码器
class Decoder(nn.Module):
    def __init__(self, word_count, output_dim, embed_dim, hidden_dim, n_layers, dropout):
        super().__init__()
        self.output_dim = output_dim
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers
        self.embedding = nn.Embedding(word_count, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, n_layers, dropout=dropout)
        self.cls = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, input, hidden, cell):
        input = input.unsqueeze(0)
        embeded = self.dropout(self.embedding(input))
        output, (hidden, cell) = self.lstm(embeded, (hidden, cell))
        prediction = self.cls(output.sequeeze(0))
        return prediction, hidden, cell
